/*
 *   Copyright 2021 zzh
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 *
 */

package com.aduib.boot.common.zookeeper;

import com.aduib.boot.common.exception.CommonException;
import com.aduib.boot.common.log.LOG;
import com.aduib.boot.common.util.StrUtils;
import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.api.CuratorWatcher;
import org.apache.curator.framework.recipes.cache.ChildData;
import org.apache.curator.framework.recipes.cache.NodeCache;
import org.apache.curator.framework.recipes.cache.NodeCacheListener;
import org.apache.curator.framework.state.ConnectionState;
import org.apache.curator.framework.state.ConnectionStateListener;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;

import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executor;
import java.util.concurrent.TimeUnit;

/**
 * @description: ZkClient
 * @author: zzh
 * @date: 2022/1/27 09:40
 */
public class ZkClient {

  private static Map<String, NodeCache> nodeCacheMap = new ConcurrentHashMap<>();

  private ZkConfig zkConfig;

  private CuratorFramework client;

  public ZkClient(ZkConfig zkConfig) {
    this.zkConfig = zkConfig;
    initClient(zkConfig);
  }

  /**
   * 静态获取ZkClient
   *
   * @param config
   * @return
   */
  public static ZkClient getClient(ZkConfig config) {
    return new ZkClient(config);
  }

  /**
   * 初始化客户端
   *
   * @param zkConfig
   */
  private void initClient(ZkConfig zkConfig) {
    if (this.zkConfig == null) {
      this.zkConfig = zkConfig;
    }
    RetryPolicy retryPolicy = zkConfig.getRetryPolicy();
    try {
      CuratorFrameworkFactory.Builder builder =
          CuratorFrameworkFactory.builder()
              .connectString(zkConfig.getAddress())
              .sessionTimeoutMs(zkConfig.getSessionTimeOut())
              .connectionTimeoutMs(zkConfig.getConnectionTimeOut())
              .retryPolicy(retryPolicy);
      String authority = zkConfig.getAuthority();
      if (StrUtils.isNotBlank(zkConfig.getAuthority())) {
        builder = builder.authorization("digest", authority.getBytes());
      }
      if (StrUtils.isNotBlank(zkConfig.getNamespace())) {
        builder = builder.namespace(zkConfig.getNamespace());
      }
      client = builder.build();
      client
          .getConnectionStateListenable()
          .addListener(new ZkClientConnectionStateListener(zkConfig));
      client.start();
      boolean connected =
          client.blockUntilConnected(zkConfig.getConnectionTimeOut(), TimeUnit.MILLISECONDS);
      if (!connected) {
        throw new CommonException("zookeeper not connected");
      }
    } catch (Exception e) {
      throw new CommonException(e.getMessage(), e);
    }
  }

  public void createPersistent(String path) {
    try {
      client.create().forPath(path);
    } catch (KeeperException.NodeExistsException e) {
      LOG.w("ZNode " + path + " already exists.", e);
    } catch (Exception e) {
      throw new IllegalStateException(e.getMessage(), e);
    }
  }

  /**
   * 创建永久有序节点
   *
   * @param path
   */
  public void createPersistentSeq(String path) {
    try {
      client.create().withMode(CreateMode.PERSISTENT_SEQUENTIAL).forPath(path);
    } catch (KeeperException.NodeExistsException e) {
      LOG.w("ZNode " + path + " already exists.", e);
    } catch (Exception e) {
      throw new IllegalStateException(e.getMessage(), e);
    }
  }

  /**
   * 创建永久有序节点-在指定时间之后没有修内容则失效
   *
   * @param path
   */
  public void createPersistentSeqTtl(String path, long ttl) {
    try {
      client
          .create()
          .withTtl(ttl)
          .withMode(CreateMode.PERSISTENT_SEQUENTIAL_WITH_TTL)
          .forPath(path);
    } catch (KeeperException.NodeExistsException e) {
      LOG.w("ZNode " + path + " already exists.", e);
    } catch (Exception e) {
      throw new IllegalStateException(e.getMessage(), e);
    }
  }

  /**
   * 创建永久节点-在指定时间之后没有修内容则失效
   *
   * @param path
   */
  public void createPersistentTtl(String path, long ttl) {
    try {
      client.create().withTtl(ttl).withMode(CreateMode.PERSISTENT_WITH_TTL).forPath(path);
    } catch (KeeperException.NodeExistsException e) {
      LOG.w("ZNode " + path + " already exists.", e);
    } catch (Exception e) {
      throw new IllegalStateException(e.getMessage(), e);
    }
  }

  public void createEphemeral(String path) {
    try {
      client.create().withMode(CreateMode.EPHEMERAL).forPath(path);
    } catch (KeeperException.NodeExistsException e) {
      LOG.w(
          "ZNode "
              + path
              + " already exists, since we will only try to recreate a node on a session expiration"
              + ", this duplication might be caused by a delete delay from the zk server, which means the old expired session"
              + " may still holds this ZNode and the server just hasn't got time to do the deletion. In this case, "
              + "we can just try to delete and create again.",
          e);
      deletePath(path);
      createEphemeral(path);
    } catch (Exception e) {
      throw new IllegalStateException(e.getMessage(), e);
    }
  }

  /**
   * 创建临时有序节点
   *
   * @param path
   */
  public void createEphemeralSeq(String path) {
    try {
      client.create().withMode(CreateMode.EPHEMERAL_SEQUENTIAL).forPath(path);
    } catch (KeeperException.NodeExistsException e) {
      LOG.w(
          "ZNode "
              + path
              + " already exists, since we will only try to recreate a node on a session expiration"
              + ", this duplication might be caused by a delete delay from the zk server, which means the old expired session"
              + " may still holds this ZNode and the server just hasn't got time to do the deletion. In this case, "
              + "we can just try to delete and create again.",
          e);
      deletePath(path);
      createEphemeral(path);
    } catch (Exception e) {
      throw new IllegalStateException(e.getMessage(), e);
    }
  }

  public void createPersistent(String path, String data) {
    byte[] dataBytes = data.getBytes(StandardCharsets.UTF_8);
    try {
      client.create().forPath(path, dataBytes);
    } catch (KeeperException.NodeExistsException e) {
      try {
        client.setData().forPath(path, dataBytes);
      } catch (Exception e1) {
        throw new IllegalStateException(e.getMessage(), e1);
      }
    } catch (Exception e) {
      throw new IllegalStateException(e.getMessage(), e);
    }
  }

  public void createPersistentSeq(String path, String data) {
    byte[] dataBytes = data.getBytes(StandardCharsets.UTF_8);
    try {
      client.create().withMode(CreateMode.PERSISTENT_SEQUENTIAL).forPath(path, dataBytes);
    } catch (KeeperException.NodeExistsException e) {
      try {
        client.setData().forPath(path, dataBytes);
      } catch (Exception e1) {
        throw new IllegalStateException(e.getMessage(), e1);
      }
    } catch (Exception e) {
      throw new IllegalStateException(e.getMessage(), e);
    }
  }

  public void createPersistentSeqTtl(String path, long ttl, String data) {
    byte[] dataBytes = data.getBytes(StandardCharsets.UTF_8);
    try {
      client
          .create()
          .withTtl(ttl)
          .withMode(CreateMode.PERSISTENT_SEQUENTIAL_WITH_TTL)
          .forPath(path, dataBytes);
    } catch (KeeperException.NodeExistsException e) {
      try {
        client.setData().forPath(path, dataBytes);
      } catch (Exception e1) {
        throw new IllegalStateException(e.getMessage(), e1);
      }
    } catch (Exception e) {
      throw new IllegalStateException(e.getMessage(), e);
    }
  }

  public void createPersistentTtl(String path, long ttl, String data) {
    byte[] dataBytes = data.getBytes(StandardCharsets.UTF_8);
    try {
      client
          .create()
          .withTtl(ttl)
          .withMode(CreateMode.PERSISTENT_WITH_TTL)
          .forPath(path, dataBytes);
    } catch (KeeperException.NodeExistsException e) {
      try {
        client.setData().forPath(path, dataBytes);
      } catch (Exception e1) {
        throw new IllegalStateException(e.getMessage(), e1);
      }
    } catch (Exception e) {
      throw new IllegalStateException(e.getMessage(), e);
    }
  }

  public void createEphemeral(String path, String data) {
    byte[] dataBytes = data.getBytes(StandardCharsets.UTF_8);
    try {
      client.create().withMode(CreateMode.EPHEMERAL).forPath(path, dataBytes);
    } catch (KeeperException.NodeExistsException e) {
      LOG.w(
          "ZNode "
              + path
              + " already exists, since we will only try to recreate a node on a session expiration"
              + ", this duplication might be caused by a delete delay from the zk server, which means the old expired session"
              + " may still holds this ZNode and the server just hasn't got time to do the deletion. In this case, "
              + "we can just try to delete and create again.",
          e);
      deletePath(path);
      createEphemeral(path, data);
    } catch (Exception e) {
      throw new IllegalStateException(e.getMessage(), e);
    }
  }

  public void deletePath(String path) {
    try {
      client.delete().deletingChildrenIfNeeded().forPath(path);
    } catch (KeeperException.NoNodeException ignored) {
    } catch (Exception e) {
      throw new IllegalStateException(e.getMessage(), e);
    }
  }

  public List<String> getChildren(String path) {
    try {
      return client.getChildren().forPath(path);
    } catch (KeeperException.NoNodeException e) {
      return null;
    } catch (Exception e) {
      throw new IllegalStateException(e.getMessage(), e);
    }
  }

  public boolean checkExists(String path) {
    try {
      if (client.checkExists().forPath(path) != null) {
        return true;
      }
    } catch (Exception e) {
    }
    return false;
  }

  public boolean isConnected() {
    return client.getZookeeperClient().isConnected();
  }

  public String getContent(String path) {
    try {
      byte[] dataBytes = client.getData().forPath(path);
      return (dataBytes == null || dataBytes.length == 0)
          ? null
          : new String(dataBytes, StandardCharsets.UTF_8);
    } catch (KeeperException.NoNodeException e) {
      // ignore NoNode Exception.
    } catch (Exception e) {
      throw new IllegalStateException(e.getMessage(), e);
    }
    return null;
  }

  public void setContent(String path, String data) {
    byte[] dataBytes = data.getBytes(StandardCharsets.UTF_8);
    try {
      client.setData().forPath(path, dataBytes);
    } catch (Exception e1) {
      throw new IllegalStateException(e1.getMessage(), e1);
    }
  }

  public void close() {
    client.close();
  }

  public CuratorWatcherImpl createTargetChildListener(
      String path, ChildListener listener) {
    return new CuratorWatcherImpl(client, listener, path);
  }

  protected List<String> addTargetChildListener(String path, CuratorWatcherImpl listener) {
    try {
      return client.getChildren().usingWatcher(listener).forPath(path);
    } catch (KeeperException.NoNodeException e) {
      return null;
    } catch (Exception e) {
      throw new IllegalStateException(e.getMessage(), e);
    }
  }

  public NodeCacheListenerImpl createTargetDataListener(String path, DataListener listener) {
    return new NodeCacheListenerImpl(client, listener, path);
  }

  public void addTargetDataListener(String path, NodeCacheListenerImpl nodeCacheListener) {
    this.addTargetDataListener(path, nodeCacheListener, null);
  }

  public void addTargetDataListener(
      String path, NodeCacheListenerImpl nodeCacheListener, Executor executor) {
    try {
      NodeCache nodeCache = new NodeCache(client, path);
      if (nodeCacheMap.putIfAbsent(path, nodeCache) != null) {
        return;
      }
      if (executor == null) {
        nodeCache.getListenable().addListener(nodeCacheListener);
      } else {
        nodeCache.getListenable().addListener(nodeCacheListener, executor);
      }

      nodeCache.start();
    } catch (Exception e) {
      throw new IllegalStateException("Add nodeCache listener for path:" + path, e);
    }
  }

  protected void removeTargetDataListener(String path, NodeCacheListenerImpl nodeCacheListener) {
    NodeCache nodeCache = nodeCacheMap.get(path);
    if (nodeCache != null) {
      nodeCache.getListenable().removeListener(nodeCacheListener);
    }
    nodeCacheListener.dataListener = null;
  }

  public void removeTargetChildListener(String path, CuratorWatcherImpl listener) {
    listener.unwatch();
  }

  static class NodeCacheListenerImpl implements NodeCacheListener {

    private CuratorFramework client;

    private volatile DataListener dataListener;

    private String path;

    protected NodeCacheListenerImpl() {}

    public NodeCacheListenerImpl(CuratorFramework client, DataListener dataListener, String path) {
      this.client = client;
      this.dataListener = dataListener;
      this.path = path;
    }

    @Override
    public void nodeChanged() throws Exception {
      ChildData childData = nodeCacheMap.get(path).getCurrentData();
      String content = null;
      Watcher.Event.EventType eventType;
      if (childData == null) {
        eventType = Watcher.Event.EventType.NodeDeleted;
      } else if (childData.getStat().getVersion() == 0) {
        content = new String(childData.getData(), StandardCharsets.UTF_8);
        eventType = Watcher.Event.EventType.NodeCreated;
      } else {
        content = new String(childData.getData(), StandardCharsets.UTF_8);
        eventType = Watcher.Event.EventType.NodeDataChanged;
      }
      dataListener.dataChanged(path, content, eventType);
    }
  }

  static class CuratorWatcherImpl implements CuratorWatcher {

    private CuratorFramework client;
    private volatile ChildListener childListener;
    private String path;

    public CuratorWatcherImpl(CuratorFramework client, ChildListener listener, String path) {
      this.client = client;
      this.childListener = listener;
      this.path = path;
    }

    protected CuratorWatcherImpl() {}

    public void unwatch() {
      this.childListener = null;
    }

    @Override
    public void process(WatchedEvent event) throws Exception {
      // if client connect or disconnect to server, zookeeper will queue
      // watched event(Watcher.Event.EventType.None, .., path = null).
      if (event.getType() == Watcher.Event.EventType.None) {
        return;
      }

      if (childListener != null) {
        childListener.childChanged(path, client.getChildren().usingWatcher(this).forPath(path));
      }
    }
  }

  private class ZkClientConnectionStateListener implements ConnectionStateListener {

    private final long UNKNOWN_SESSION_ID = -1L;

    private long lastSessionId;

    private int timeout;

    private int sessionExpireMs;

    public ZkClientConnectionStateListener(ZkConfig zkConfig) {
      this.timeout = zkConfig.getConnectionTimeOut();
      this.sessionExpireMs = zkConfig.getSessionTimeOut();
    }

    @Override
    public void stateChanged(CuratorFramework client, ConnectionState state) {
      long sessionId = UNKNOWN_SESSION_ID;
      try {
        sessionId = client.getZookeeperClient().getZooKeeper().getSessionId();
      } catch (Exception e) {
        LOG.w("Zk client state changed, but failed to get the related zk session instance.");
      }

      if (state == ConnectionState.LOST) {
        LOG.w("zookeeper session " + Long.toHexString(lastSessionId) + " expired.");
      } else if (state == ConnectionState.SUSPENDED) {
        LOG.w(
            "zookeeper connection of session "
                + Long.toHexString(sessionId)
                + " timed out. "
                + "connection timeout value is "
                + timeout
                + ", session expire timeout value is "
                + sessionExpireMs);
      } else if (state == ConnectionState.CONNECTED) {
        lastSessionId = sessionId;
        LOG.i(
            "zookeeper client instance initiated successfully, session id is "
                + Long.toHexString(sessionId));
      } else if (state == ConnectionState.RECONNECTED) {
        if (lastSessionId == sessionId && sessionId != UNKNOWN_SESSION_ID) {
          LOG.w(
              "zookeeper connection recovered from connection lose, "
                  + "reuse the old session "
                  + Long.toHexString(sessionId));
        } else {
          LOG.w(
              "New session created after old session lost, "
                  + "old session "
                  + Long.toHexString(lastSessionId)
                  + ", new session "
                  + Long.toHexString(sessionId));
          lastSessionId = sessionId;
        }
      }
    }
  }
}
